## ---------------------------------------------------------------------- #
## Figure 1 and Figure 2
## Christopher J. Fariss
## "Yes, Human Rights Practices Are Improving Over Time"
## American Political Science Review
## ---------------------------------------------------------------------- #
rm(list = ls())

## ---------------------------------------------------------------------- #
## load packages
pkgs <- c("rjags", "rstan", "runjags", "parallel", "reshape")
invisible(sapply(pkgs, require, character.only = TRUE))

## set seed for replication
set.seed(12345)

## simulate data
N <- 30 # units
T <- 10 # time periods
N_T <- N*T
x <- matrix(NA,nrow=N,ncol=T)

x[,1] <- rnorm(N,0,1) # random draw for first value
for(t in 2:T){
    x[,t] <- rnorm(N,x[,t-1],0.05) # random draw for subsequent time periods
}

SEQ <- seq(-2,2,length.out=T)
SEQ

for(t in 1:T){
    x[,t] <- SEQ[t] + x[,t]
}
apply(x, 2, mean)

Truth <- apply(x, 2, mean)

x <- melt(x)


alpha <- c(0,0,0,0,0)
beta <- c(1,1,1,1,1)
k <- 5

y <- matrix(NA, nrow=N_T, ncol=k)

for(j in 1:k){
    # linear terms of the model
    xb <- alpha[j] + beta[j] * x$value
    
    # transform the linear xb terms using the logit function
    # so that theta is bound from 0 to 1
    theta <- 1 / (1 + exp(-xb))
    
    # generate the items with theta and measurment error
    y[,j] <- rbinom(N*T, size=1, prob=theta)
}

J <- ncol(y)

subjects <- x[,1]
time <- x[,2]




# ---------------------------------------------------------------------- #
# Define JAGS model statements for 6 models

MODEL5 <- "

model{
    for(i in 1:N_T){
        
        for(j in 1:5){
            logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
            }
        #for(j in 5:5){
        #    logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
        #    y[i,j] ~ dbern(p[i,j])
        #}
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL5, file="MODEL5.bug")


MODEL4 <- "

model{
    for(i in 1:N_T){
        
        for(j in 1:4){
            logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        for(j in 5:5){
            logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL4, file="MODEL4.bug")


MODEL3 <- "

model{
    for(i in 1:N_T){
        
        for(j in 1:3){
            logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        for(j in 4:5){
            logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL3, file="MODEL3.bug")


MODEL2 <- "

model{
    for(i in 1:N_T){
        
        for(j in 1:2){
            logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        for(j in 3:5){
            logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL2, file="MODEL2.bug")


MODEL1 <- "

model{
    for(i in 1:N_T){
        
        for(j in 1:1){
            logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        for(j in 2:5){
            logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL1, file="MODEL1.bug")


MODEL0 <- "

model{
    for(i in 1:N_T){
        
        #for(j in 1:4){
        #    logit(p[i,j]) <- alpha[j,1] + beta[j] * x[i]
        #    y[i,j] ~ dbern(p[i,j])
        #}
        for(j in 1:5){
            logit(p[i,j]) <- alpha[j,time[i]] + beta[j] * x[i]
            y[i,j] ~ dbern(p[i,j])
        }
        
        x[i] <- mu[subjects[i], time[i]]
    }
    
    sigma ~ dunif(0,1)
    kappa <- pow(sigma, -1)
    
    
    for(c in 1:N){
        mu[c, 1] ~ dnorm(0, 1)
        for(t in 2:T){ #n.year is number of years
            mu[c, t] ~ dnorm(mu[c, t-1], kappa)
            #mu[c, t] ~ dnorm(0, 1)
        }
    }
    
    #x.bar  <- mean(mu)
    
    for(j in 1:J){
        for(t in 1:T){
            alpha[j,t] ~ dnorm(0,.1)
        }
        #alpha[j] ~ dnorm(0,.1)
        beta[j] ~ dnorm(0,1) T(0,)
    }
}
"

# ---------------------------------------------------------------------- #
# write the file as a temporary name to then read in
write(MODEL0, file="MODEL0.bug")



for(k in 0:5){
# ---------------------------------------------------------------------- #
# create initial values for the latent variable model

MakeInits <- function(){
    list(mu=matrix(rnorm(N_T), nrow=N,ncol=T),
    #alpha=rnorm(J),
    alpha=array(rnorm(J*T),dim=c(J,T)),
    beta=runif(J),
    sigma=runif(1)
    )
}

inits.function <- function(chain){
    return(switch(chain,
    "1"=MakeInits(),
    "2"=MakeInits(),
    "3"=MakeInits(),
    "4"=MakeInits(),
    "5"=MakeInits(),
    "6"=MakeInits(),
    "7"=MakeInits(),
    "8"=MakeInits()
    ))
}

# ---------------------------------------------------------------------- #
# set variables for JAGS model call
CHAINS <- 4
ADAPT <- 1000
BURNIN <- 2000
DRAWS <- 2000
THIN <- 10
#CLUSTER <- makeCluster(4)

# set model file for JAGS model call
if(k==5) MODEL.FILE <- "MODEL5.bug"
if(k==4) MODEL.FILE <- "MODEL4.bug"
if(k==3) MODEL.FILE <- "MODEL3.bug"
if(k==2) MODEL.FILE <- "MODEL2.bug"
if(k==1) MODEL.FILE <- "MODEL1.bug"
if(k==0) MODEL.FILE <- "MODEL0.bug"

# ---------------------------------------------------------------------- #
# runjags version of call to JAGS

load.module("glm")
list.modules()


m <- jags.model(file=MODEL.FILE,
    data = list("y"=y, "N"=N, "T"=T, "N_T"=N_T, "J"=J, "subjects"=subjects, "time"=time),
    n.chains = CHAINS,
    inits = inits.function,
    n.adapt = ADAPT
)

update(m, BURNIN)
M <- coda.samples(m, DRAWS, variable.names=c("x"), THIN)



# ---------------------------------------------------------------------- #
# process JAGS estimates
mat1 <- as.matrix(as.mcmc(M[[1]]))
mat2 <- as.matrix(as.mcmc(M[[2]]))
mat3 <- as.matrix(as.mcmc(M[[3]]))
mat4 <- as.matrix(as.mcmc(M[[4]]))

posterior_estimates <- rbind(mat1, mat2, mat3, mat4)
parameter.mean <- apply(posterior_estimates, 2, mean)
parameter.sd <- apply(posterior_estimates, 2, sd)

posterior.mu <- mean(parameter.mean)

year.mean <- year.lb <- year.ub <- NA
for(i in 1:max(time)){
    year.mean[i] <- mean(apply(posterior_estimates[,time==i],1,mean)) - posterior.mu
    year.lb[i] <- quantile(apply(posterior_estimates[,time==i],1,mean), .025) - posterior.mu
    year.ub[i] <- quantile(apply(posterior_estimates[,time==i],1,mean), .975) - posterior.mu
}

# plot true latent variable with posterior mean
par(mar=c(4,4,1,1), font=2, font.lab=2, cex=1.3)
plot(1:T, year.mean, ylim=c(-3,3), ylab="Latent Variable Average Over Time", xlab="Time", type="n")

for(i in 1:T){
    lines(c(i,i), c(year.lb[i], year.ub[i]), lwd=2, col=grey(.6))
}
points(1:T, year.mean, col=grey(.6), pch=21, bg=grey(.95), cex=.8)

points(1:T, Truth, col=1, pch=19)
#points(1:T, c(-0.1326729,  0.1118860,  0.3571317,  0.5982520,  0.8379022 , 1.0917770,  1.3420348,  1.5933299,  1.8499121,  2.0903428), col=2)


value <- NA
for(i in 1:800){value[i] <- cor(posterior_estimates[i,], x$value, method="spearman")}
summary(value)


if(k==0) prob0 <- as.numeric(quantile(value, c(.025,.5,.975)))
if(k==1) prob1 <- as.numeric(quantile(value, c(.025,.5,.975)))
if(k==2) prob2 <- as.numeric(quantile(value, c(.025,.5,.975)))
if(k==3) prob3 <- as.numeric(quantile(value, c(.025,.5,.975)))
if(k==4) prob4 <- as.numeric(quantile(value, c(.025,.5,.975)))
if(k==5) prob5 <- as.numeric(quantile(value, c(.025,.5,.975)))


}


par(mar=c(4,5,1,1))
barplot(c(prob5[2], prob4[2], prob3[2], prob2[2], prob1[2], prob0[2]), space=0, ylim=c(0,1), col=grey(.85), yaxt="n", names.arg=c("5", "4", "3", "2", "1", "0"))
#mtext(side=3, "Relationship between Estimated Latent Trait and True Trait", line=1)
mtext(side=2, "Correlation Coefficient \nbetween Estimated Latent Trait and Simulated Trait", line=3)
#mtext(side=1, "Number of Varying Difficulty Parameters", line=3)
mtext(side=1, "Number of Fixed (Constant) Difficulty Parameters", line=3)

axis(side=2, at=c(0,.2,.4,.6,.8,1),las=2)
lines(c(0.5,0.5), c(prob5[1],prob5[3]))
lines(c(1.5,1.5), c(prob4[1],prob4[3]))
lines(c(2.5,2.5), c(prob3[1],prob3[3]))
lines(c(3.5,3.5), c(prob2[1],prob2[3]))
lines(c(4.5,4.5), c(prob1[1],prob1[3]))
lines(c(5.5,5.5), c(prob0[1],prob0[3]))
